import numpy as np
from pykalman import KalmanFilter
from KalmanFilter import VehicleKalmanFilter
class DriverScoringSystem:
    def __init__(self):
        # 特征提取参数
        self.window_size = 10  # 滑动窗口大小
        self.acceleration_threshold = 2.5  # 急加速阈值 (m/s²)
        self.brake_threshold = -3.0  # 急刹车阈值 (m/s²)
        self.turn_threshold = 1.2  # 急转弯阈值 (rad/s)
        self.jerk_threshold = 15.0  # 加加速度阈值 (m/s³)

        # 颠簸阈值
        self.bump_thresholds = {
            'longitudinal': {
                'slight': (1.0, 1.5),
                'obvious': (1.5, 2.5),
                'dangerous': 3.0
            },
            'vertical': {
                'slight': (0.5, 1.0),
                'obvious': (1.0, 2.0),
                'dangerous': 2.5
            },
            'lateral': {
                'slight': (0.5, 0.8),
                'obvious': (0.8, 1.2),
                'dangerous': 1.5
            }
        }

        # 数据缓冲区
        self.acceleration_buffer = []
        self.orientation_buffer = []
        self.time_buffer = []

        # 事件计数器
        self.hard_acceleration_count = 0
        self.hard_brake_count = 0
        self.sharp_turn_count = 0
        self.rough_road_count = 0
        self.slight_bump_count = 0
        self.obvious_bump_count = 0
        self.dangerous_bump_count = 0

        # 针对ICM42688的卡尔曼滤波器
        self.kalman_filter = VehicleKalmanFilter()
        self.filter_state = None
        self.filter_covariance = None

    def _init_kalman_filter(self):
        """初始化卡尔曼滤波器"""
        # 状态转移矩阵 (姿态预测模型)
        transition_matrix = np.array([
            [1, 0, 0, 1, 0, 0],  # 位置x = 位置x + 速度x
            [0, 1, 0, 0, 1, 0],  # 位置y = 位置y + 速度y
            [0, 0, 1, 0, 0, 1],  # 位置z = 位置z + 速度z
            [0, 0, 0, 1, 0, 0],  # 速度x = 速度x
            [0, 0, 0, 0, 1, 0],  # 速度y = 速度y
            [0, 0, 0, 0, 0, 1]   # 速度z = 速度z
        ])

        # 观测矩阵 (从状态到观测的映射)
        observation_matrix = np.array([
            [1, 0, 0, 0, 0, 0],  # 观测位置x
            [0, 1, 0, 0, 0, 0],  # 观测位置y
            [0, 0, 1, 0, 0, 0]   # 观测位置z
        ])

        # 状态转移噪声协方差
        transition_covariance = 1e-4 * np.eye(6)

        # 观测噪声协方差
        observation_covariance = 1e-2 * np.eye(3)

        # 初始状态协方差
        initial_covariance = np.eye(6)

        # 初始化卡尔曼滤波器
        kf = KalmanFilter(
            transition_matrices=transition_matrix,
            observation_matrices=observation_matrix,
            transition_covariance=transition_covariance,
            observation_covariance=observation_covariance,
            initial_state_covariance=initial_covariance
        )

        return kf

    def update_buffers(self, acceleration, orientation, timestamp):
        """更新滑动窗口缓冲区"""
        self.acceleration_buffer.append(acceleration)
        self.orientation_buffer.append(orientation)
        self.time_buffer.append(timestamp)

        # 保持缓冲区固定大小
        if len(self.acceleration_buffer) > self.window_size:
            self.acceleration_buffer.pop(0)
            self.orientation_buffer.pop(0)
            self.time_buffer.pop(0)

    def calculate_jerk(self):
        """计算加加速度(jerk) - 加速度的变化率"""
        if len(self.acceleration_buffer) < 3:
            return np.zeros(3)

        # 获取最近3个加速度样本
        a1, a2, a3 = self.acceleration_buffer[-3:]
        t1, t2, t3 = self.time_buffer[-3:]

        # 计算平均加加速度
        dt1 = t2 - t1
        dt2 = t3 - t2

        if dt1 <= 0 or dt2 <= 0:
            return np.zeros(3)

        jerk1 = (a2 - a1) / dt1
        jerk2 = (a3 - a2) / dt2

        return (jerk1 + jerk2) / 2

    def detect_bumps(self, acceleration):
        """检测不同类型的颠簸事件"""
        longitudinal_acc = acceleration[0]
        vertical_acc = acceleration[1]
        lateral_acc = acceleration[2]

        slight_bump = False
        obvious_bump = False
        dangerous_bump = False

        # 纵向颠簸检测
        if (self.bump_thresholds['longitudinal']['slight'][0] <= abs(longitudinal_acc) < self.bump_thresholds['longitudinal']['slight'][1] or
            self.bump_thresholds['longitudinal']['obvious'][0] <= abs(longitudinal_acc) < self.bump_thresholds['longitudinal']['obvious'][1] or
            abs(longitudinal_acc) >= self.bump_thresholds['longitudinal']['dangerous']):
            if self.bump_thresholds['longitudinal']['slight'][0] <= abs(longitudinal_acc) < self.bump_thresholds['longitudinal']['slight'][1]:
                slight_bump = True
            elif self.bump_thresholds['longitudinal']['obvious'][0] <= abs(longitudinal_acc) < self.bump_thresholds['longitudinal']['obvious'][1]:
                obvious_bump = True
            else:
                dangerous_bump = True

        # 垂直颠簸检测
        if (self.bump_thresholds['vertical']['slight'][0] <= abs(vertical_acc) < self.bump_thresholds['vertical']['slight'][1] or
            self.bump_thresholds['vertical']['obvious'][0] <= abs(vertical_acc) < self.bump_thresholds['vertical']['obvious'][1] or
            abs(vertical_acc) >= self.bump_thresholds['vertical']['dangerous']):
            if self.bump_thresholds['vertical']['slight'][0] <= abs(vertical_acc) < self.bump_thresholds['vertical']['slight'][1]:
                slight_bump = True
            elif self.bump_thresholds['vertical']['obvious'][0] <= abs(vertical_acc) < self.bump_thresholds['vertical']['obvious'][1]:
                obvious_bump = True
            else:
                dangerous_bump = True

        # 横向颠簸检测
        if (self.bump_thresholds['lateral']['slight'][0] <= abs(lateral_acc) < self.bump_thresholds['lateral']['slight'][1] or
            self.bump_thresholds['lateral']['obvious'][0] <= abs(lateral_acc) < self.bump_thresholds['lateral']['obvious'][1] or
            abs(lateral_acc) >= self.bump_thresholds['lateral']['dangerous']):
            if self.bump_thresholds['lateral']['slight'][0] <= abs(lateral_acc) < self.bump_thresholds['lateral']['slight'][1]:
                slight_bump = True
            elif self.bump_thresholds['lateral']['obvious'][0] <= abs(lateral_acc) < self.bump_thresholds['lateral']['obvious'][1]:
                obvious_bump = True
            else:
                dangerous_bump = True

        if slight_bump:
            self.slight_bump_count += 1
        if obvious_bump:
            self.obvious_bump_count += 1
        if dangerous_bump:
            self.dangerous_bump_count += 1

        return slight_bump, obvious_bump, dangerous_bump

    def kalman_filter_update(self, gyroscope, acceleration, dt):
        """使用新的ICM42688卡尔曼滤波器更新"""
        # 调用新的卡尔曼滤波器
        filtered_state = self.kalman_filter.update(acceleration, gyroscope, dt)

        # 从滤波后的状态中提取姿态信息
        pitch, roll, yaw, _, _, _ = filtered_state

        # 返回滤波后的姿态
        return np.array([pitch, roll, yaw])

    def extract_features(self, timestamp):
        """提取驾驶行为特征"""
        if len(self.acceleration_buffer) < self.window_size:
            return {}

        # 获取最新的加速度和姿态数据
        latest_acceleration = self.acceleration_buffer[-1]
        latest_orientation = self.orientation_buffer[-1]

        # 计算时间间隔
        if len(self.time_buffer) >= 2:
            dt = self.time_buffer[-1] - self.time_buffer[-2]
        else:
            dt = 0.1  # 默认时间间隔

        # 应用卡尔曼滤波
        filtered_orientation = self.kalman_filter_update(latest_orientation, latest_acceleration, dt)

        # 1. 计算纵向加速度 (假设X轴为前进方向)
        longitudinal_acc = latest_acceleration[0]

        # 2. 计算转弯速率 (使用滤波后的姿态)
        turn_rate = np.abs(filtered_orientation[2])

        # 3. 计算加加速度
        jerk = self.calculate_jerk()
        jerk_magnitude = np.linalg.norm(jerk)

        # 4. 检测驾驶事件
        is_hard_acceleration = longitudinal_acc > self.acceleration_threshold
        is_hard_brake = longitudinal_acc < self.brake_threshold
        is_sharp_turn = turn_rate > self.turn_threshold
        is_rough_road = jerk_magnitude > self.jerk_threshold

        # 5. 检测颠簸事件
        slight_bump, obvious_bump, dangerous_bump = self.detect_bumps(latest_acceleration)

        # 6. 更新事件计数器
        if is_hard_acceleration:
            self.hard_acceleration_count += 1
        if is_hard_brake:
            self.hard_brake_count += 1
        if is_sharp_turn:
            self.sharp_turn_count += 1
        if is_rough_road:
            self.rough_road_count += 1

        # 返回提取的特征
        event = {
            'longitudinal_acceleration': longitudinal_acc,
            'turn_rate': turn_rate,
            'jerk_magnitude': jerk_magnitude,
            'is_hard_acceleration': is_hard_acceleration,
            'is_hard_brake': is_hard_brake,
            'is_sharp_turn': is_sharp_turn,
            'is_rough_road': is_rough_road,
            'is_slight_bump': slight_bump,
            'is_obvious_bump': obvious_bump,
            'is_dangerous_bump': dangerous_bump,
            # 确保所有可能的键都存在
            'hard_acceleration': longitudinal_acc if is_hard_acceleration else 0,
            'hard_brake': longitudinal_acc if is_hard_brake else 0,
            'sharp_turn': turn_rate if is_sharp_turn else 0,
            'rough_road': jerk_magnitude if is_rough_road else 0,
            'slight_bump': latest_acceleration if slight_bump else 0,
            'obvious_bump': latest_acceleration if obvious_bump else 0,
            'dangerous_bump': latest_acceleration if dangerous_bump else 0
        }
        return event

    def get_event_metrics(self):
        """获取累积的事件统计数据"""
        return {
            'hard_acceleration_count': self.hard_acceleration_count,
            'hard_brake_count': self.hard_brake_count,
            'sharp_turn_count': self.sharp_turn_count,
            'rough_road_count': self.rough_road_count,
            'slight_bump_count': self.slight_bump_count,
            'obvious_bump_count': self.obvious_bump_count,
            'dangerous_bump_count': self.dangerous_bump_count
        }

    import numpy as np
    from pykalman import KalmanFilter


        # 原有的代码保持不变

    def get_filtered_data(self):
        """获取滤波后的加速度、姿态、角速度数据"""
        if len(self.acceleration_buffer) < self.window_size:
            return None, None, None

        # 获取最新的加速度和姿态数据
        latest_acceleration = self.acceleration_buffer[-1]
        latest_orientation = self.orientation_buffer[-1]

        # 计算时间间隔
        if len(self.time_buffer) >= 2:
            dt = self.time_buffer[-1] - self.time_buffer[-2]
        else:
            dt = 0.1  # 默认时间间隔

        # 应用卡尔曼滤波
        filtered_state = self.kalman_filter.update(latest_acceleration, latest_orientation, dt)

        # 从滤波后的状态中提取姿态信息
        filtered_orientation = filtered_state[:3]

        # 获取滤波后的加速度
        filtered_acceleration = self.kalman_filter.get_filtered_acceleration(latest_acceleration)

        print(f"Raw orientation: {latest_orientation}, Filtered orientation: {filtered_orientation}")
        print(f"Raw acceleration: {latest_acceleration}, Filtered acceleration: {filtered_acceleration}")

        return filtered_acceleration, filtered_orientation, latest_orientation
